#include <stdio.h>
#include <string.h>
#include <math.h>
#include <stdlib.h>
#include <sys/time.h>
#include <time.h>
#include <mpi.h>

#define EPSILON 1.0e-5f
#define REAL_T float
#define C 0.2f

#ifdef _OPENMP
#include <omp.h>
#endif

#define TAG_FIRST_LINE 42
#define TAG_LAST_LINE 43

/* Pour obtenir le temps "horloge" : */
double my_gettimeofday(){
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return tv.tv_sec + (tv.tv_usec * 1.0e-6L);
}

static inline REAL_T calcul_iter_no_tests(REAL_T **grid[2], int grid_nb, int i, int j) {
	return ((1.0f - 4.0f * C) * grid[grid_nb][i][j]
	  + C * (grid[grid_nb][i][j-1])
	  + C * (grid[grid_nb][i][j+1])
	  + C * (grid[grid_nb][i-1][j]) 
	  + C * (grid[grid_nb][i+1][j]));
}

static inline REAL_T calcul_iter_tests (REAL_T **grid[2], int grid_nb, int i, int j, int h, int w, int first_index) {
	return ((1.0f - 4.0f * C) * grid[grid_nb][i][j]
	  + C * (j>0   ? grid[grid_nb][i][j-1] : grid[grid_nb][i][j])
	  + C * (j<w-1 ? grid[grid_nb][i][j+1] : grid[grid_nb][i][j])
	  + C * (i>0   ? grid[grid_nb][i-1][j] : grid[grid_nb][i][j]) 
	  + C * (i+first_index<h-1 ? grid[grid_nb][i+1][j] : grid[grid_nb][i][j]));
}

int main(int argc, char **argv){

    /******************/
    /* INITIALISATION */
    /******************/
    
    int n = 0;
    int w = 0;
    int h = 0;
    int i = 0, j = 0;
    int step = 0;

    int grid_nb = 0, inv_grid_nb=0;
    REAL_T **grid[2] = {NULL, NULL};

    REAL_T max_error = 0.0f;

    double t_debut = 0.0, t_fin = 0.0;
    double t_total_debut = 0.0, t_total_fin = 0.0;

    /* Variables MPI */
    int rank;
    int size;
    int h_local, first_index;
    int iter;
    REAL_T max_error_local = 0.0f;
    int h_debut, h_fin;
    MPI_Status status;
    MPI_Request req_first, req_last, req_r_first, req_r_last;

    /* Valeurs par defaut : */
    w = 3;
    h = 500;

    /* Parametres : */
    if( argc > 1) h    = atoi(argv[1]);
    if( argc > 2) w    = atoi(argv[2]);
	 
    /* Initialisation MPI */
    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);    
    
    first_index = h/size * rank - (rank>0);
    h_local = h/size + (rank>0) + (rank<size-1);
    
    h_debut = (rank>0);
    h_fin = h_local - (rank < (size-1));
    
    
   /* Ne fonctionne que pour w >= 3 et h_local >= 3 */
   if(w<3 || h_local<3) {
		printf("w ou h_local doit être >= à 3\n");
		MPI_Finalize();
		exit(1);
	}
    
    if(rank == 0) {
        /* Demarrage horloge : */
        t_total_debut = my_gettimeofday();

        /* Affichage informations diverses : */
        printf("Dims : (%d, %d) \t Nb points : %i \t Taille memoire : %f Mo X2\n",
         w, h, w*h, (float) w*h*sizeof(REAL_T) / (1024*1024));
    }
    
   /* Allocation memoire : h_local */
   for (grid_nb=0; grid_nb<=1; grid_nb++){
   	grid[grid_nb] = (REAL_T **) malloc(h_local*sizeof(REAL_T *));
      if (grid[grid_nb] == NULL){
         fprintf(stderr, "Echec allocation memoire pour 'grid[%d]'.\n", grid_nb);
      }
      for (i=0; i<h_local; i++){
      	grid[grid_nb][i] = (REAL_T *) malloc(w*sizeof(REAL_T));
         if (grid[grid_nb][i] == NULL){
         	fprintf(stderr, "Echec allocation memoire pour 'grid[%d][%d]'.\n", grid_nb, i);
         }

         /* set to 0: */
      	memset(grid[grid_nb][i], 0, w*sizeof(REAL_T));
		} /* for i */
   } /* for grid_nb */

   /* Initialisation de la grille : */
   grid_nb=0;
	
	#pragma omp parallel
	{
		#pragma omp for private(j) nowait
   	for (i=0; i<h_local; i++){
     		for (j=0; j<w; j++){
         	if((i+first_index)%100 == 50)
            	grid[grid_nb][i][j] = (REAL_T) 100.0f;
         	else
            	grid[grid_nb][i][j] = (REAL_T) 25.0f;
   		}
   	}
   }

	/* Test Exhaustif */
	/*FILE * fi;
	char filenamei[10];
	sprintf(filenamei, "results_bande_p%d_init", rank);
	fi = fopen(filenamei, "w");

	for(i=h_debut;i<h_fin;i++) {
		for(j=0;j<w;j++) {
			fprintf(fi, "%2.3f ", grid[inv_grid_nb][i][j]);
		}
		fprintf(fi, "\n");
	}
	
	fclose(fi);*/

	/***** Calcul : *****/
	iter = 0;
    
	do {
    	inv_grid_nb = 1-grid_nb;

      if(rank == 0) {
  	   	/* Debut horloge pour cette iteration : */
         t_debut = my_gettimeofday();
      }
        
      /* mise a zero de grid[inv_grid_nb]: */
      for (i=0; i<h_local; i++){
      	memset(grid[inv_grid_nb][i], 0, w*sizeof(REAL_T));    
      } /* for i */
		        
		#pragma omp parallel
		{
			#pragma omp sections private(i,j)
			{
			#pragma omp section
			{
	      	/* Receive de la ligne 0 de (rank-1) -> LAST_LINE */
   	   	if(rank>0) {
      	   	MPI_Irecv(&grid[inv_grid_nb][0][0], w, MPI_FLOAT, rank-1, TAG_LAST_LINE, MPI_COMM_WORLD, &req_r_last);
      		}
      	}
        
         #pragma omp section 
			{
	     	 	/* Receive de la ligne h_local-1 de (rank+1) -> FIRST_LINE */
	      	if(rank <(size-1)) {
   	      	MPI_Irecv(&grid[inv_grid_nb][h_local-1][0], w, MPI_FLOAT, rank+1, TAG_FIRST_LINE, MPI_COMM_WORLD, &req_r_first);
     			}
     		}

			#pragma omp section
			{
				/* Calcul 1ère ligne */
				for (j=0; j<w; j++){
					i = h_debut;
      			grid[inv_grid_nb][i][j] = calcul_iter_tests(grid, grid_nb, i, j, h, w, first_index); 

		   	} /* for j */
        
   			/* Envoi de la ligne 1 à (rank-1) -> FIRST_LINE */
      		if(rank>0) {
       			MPI_Isend(&grid[inv_grid_nb][1][0], w, MPI_FLOAT, rank-1, TAG_FIRST_LINE, MPI_COMM_WORLD, &req_first);
      		}
      	}
		
			#pragma omp section
			{
				/* Calcul dernière ligne*/
				for (j=0; j<w; j++){
					i = h_fin-1;
      	   	grid[inv_grid_nb][i][j] = calcul_iter_tests(grid, grid_nb, i, j, h, w, first_index); 
      		} /* for j */
			
				/* Envoi de la ligne h_local-2 à (rank+1) -> LAST_LINE */
      		if(rank <(size-1)) {
      			MPI_Isend(&grid[inv_grid_nb][h_local-2][0], w, MPI_FLOAT, rank+1, TAG_LAST_LINE, MPI_COMM_WORLD, &req_last);
      		}
      	}
			}
			
			/* Calcul reste */
			#pragma omp for private(j)
	      for (i=h_debut+1; i<h_fin-1; i++){
      
      		grid[inv_grid_nb][i][0] = calcul_iter_tests(grid, grid_nb, i, 0, h, w, first_index); 
      
				for (j=1; j<w-1; j++){
         		grid[inv_grid_nb][i][j] = calcul_iter_no_tests(grid, grid_nb, i, j); 

         	} /* for j */
         
         	grid[inv_grid_nb][i][w-1] = calcul_iter_tests(grid, grid_nb, i, 0, h, w, first_index); 
                        
      	} /* for i */
        
        
			/* calcul de l'erreur */
			#pragma omp single
			{
				max_error_local = 0.0f;
      	}
		
      	REAL_T max_error_thread = 0.0f;
      	REAL_T error = 0.0f;
      	
      	#pragma omp for private(j) nowait
      	for (i=h_debut; i<h_fin; i++){
      		for (j=0; j<w; j++){
        	 		error = fabs(grid[grid_nb][i][j] - grid[inv_grid_nb][i][j]);
            	if (error > max_error_thread){ 
      				max_error_thread = error; 
            	}
         	} /* for j */
      	} /* for i */
      
      	#pragma omp critical
      	{
      		if (max_error_local < max_error_thread)
      			max_error_local = max_error_thread;
      	}
      }		

      /* Réduction sur le max des max_error_local */
      MPI_Allreduce(&max_error_local, &max_error, 1, MPI_FLOAT, MPI_MAX, MPI_COMM_WORLD);
        
      if(rank == 0) {
         /* Arret horloge pour cette iteration : */
         t_fin = my_gettimeofday();

         /* Affichage pour cette iteration : */
         printf("Iteration %d : \t delta = %.3e \t temps = %.1f s\n", 
             ++n, max_error, t_fin - t_debut);
      }
      
      #pragma omp parallel
      {
        
      	/* Attente fin réception (et envoi)
       	* avant d'échanger les grilles 
       	* et passer au tour suivant */
      	#pragma omp sections
      	{
      		#pragma omp section
	      	if(rank>0)
        			MPI_Wait(&req_r_last, &status);
        		
        		#pragma omp section
        		if(rank>0)
        			MPI_Wait(&req_first, &status);
      		
	    	   #pragma omp section
      		if(rank <(size-1))
        			MPI_Wait(&req_r_first, &status);
        		
        		#pragma omp section
        		if(rank <(size-1))
        			MPI_Wait(&req_last, &status);
      	}
      }
        
      /* update 'grid_nb': */
      grid_nb = inv_grid_nb;
      
      /* itération suivante */
      iter++;
	} while (max_error > EPSILON && iter < 40);
   /***** Fin du calcul *****/

	/* Test Exhaustif - Enregistrement des résultats de (1-grid_nb) */
	FILE * f;
	char filename[10];
	sprintf(filename, "results_bande_p%d", rank);
	f = fopen(filename, "w");

	for(i=h_debut;i<h_fin;i++) {
		for(j=0;j<w;j++) {
			fprintf(f, "%3.5f ", grid[inv_grid_nb][i][j]);
		}
		fprintf(f, "\n");
	}
	
	fclose(f);
	

   /* Liberation memoire : */
   for (grid_nb=0; grid_nb<=1; grid_nb++){
    
      for (i=0; i<h_local; i++){
         free(grid[grid_nb][i]);
      }
        
      free(grid[grid_nb]);
   }

	if(rank == 0) {
   	/* Arret horloge : */
      t_total_fin = my_gettimeofday();  

      /* Affichage des resultats : */
      printf("\nNombre d'iterations : %d\n", n); 
      printf("Erreur (delta) = %.3e\n", max_error); 
      printf("Temps total : %.1f s\n", t_total_fin - t_total_debut); 
	}
    
   MPI_Finalize();

   return EXIT_SUCCESS;
}


